# mysql_db.py
import os
import pymysql
from dotenv import load_dotenv
from loguru import logger

load_dotenv()  # 加载项目根目录下的 .env 文件

class MySQLDB:
    def __init__(self):
        self.host = os.getenv('MYSQL_HOST', 'localhost')
        self.port = int(os.getenv('MYSQL_PORT', 3306))
        self.user = os.getenv('MYSQL_USERNAME', 'root')
        self.password = os.getenv('MYSQL_PASSWORD', '')
        self.database = os.getenv('MYSQL_DATABASE', 'test_db')
        self.charset = os.getenv('MYSQL_CHARSET', 'utf8mb4')
        self.conn = None
        self.cursor = None

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def connect(self):
        try:
            self.conn = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                charset=self.charset,
            )
            self.cursor = self.conn.cursor()
        except Exception as e:
            logger.error(f"[×] 数据库连接失败: {e}")
            raise

    def close(self):
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()

    # 查询
    def query(self, sql, params=None):
        try:
            self.cursor.execute(sql, params)
            return self.cursor.fetchall()
        except Exception as e:
            logger.error(f"[×] 查询失败: {e}, SQL: {sql}")
            return None

    # 执行插入/更新/删除
    def execute(self, sql, params=None, commit=True):
        try:
            self.cursor.execute(sql, params)
            if commit:
                self.conn.commit()
            return self.cursor.rowcount
        except Exception as e:
            self.conn.rollback()
            logger.error(f"[×] 执行失败: {e}, SQL: {sql}")
            return 0

    # 批量执行
    def executemany(self, sql, data_list, commit=True):
        try:
            self.cursor.executemany(sql, data_list)
            if commit:
                self.conn.commit()
            return self.cursor.rowcount
        except Exception as e:
            self.conn.rollback()
            logger.error(f"[×] 批量执行失败: {e}, SQL: {sql}")
            return 0
